import torch
import torch.nn as nn
import pandas as pd
from flask import Flask, request, jsonify
from transformers import RobertaTokenizer, RobertaModel
from torch.utils.data import Dataset, DataLoader
import os
from datetime import datetime
import numpy as np
import pickle
from features import extract_domain_features
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("api.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)
from flask_cors import CORS
# 确保 model 目录存在
def ensure_model_dir():
    model_dir = "model"
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    return model_dir

# 加载 RoBERTa 预训练模型
def load_roberta(model_name="roberta-base"):
    return RobertaTokenizer.from_pretrained(model_name, mirror="tuna"), RobertaModel.from_pretrained(model_name, mirror="tuna")

# 域名分类器
class DomainAPTClassifier(nn.Module):
    def __init__(self, num_classes):
        super(DomainAPTClassifier, self).__init__()
        _, self.roberta = load_roberta()
        self.fc = nn.Linear(self.roberta.config.hidden_size, num_classes)
        self.dropout = nn.Dropout(0.5)

    def forward(self, input_ids, attention_mask):
        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0, :]
        cls_output = self.dropout(cls_output)
        logits = self.fc(cls_output)
        return logits

# 自定义数据集
class DomainDataset(Dataset):
    def __init__(self, domains, tokenizer, max_length=32):
        self.data = pd.DataFrame({"domain": domains})
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        domain = self.data.iloc[idx, 0]
        encoding = self.tokenizer(domain, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='pt')
        return {
            'input_ids': encoding['input_ids'].squeeze(0),
            'attention_mask': encoding['attention_mask'].squeeze(0)
        }

# 预测函数
def predict(model, domains, tokenizer, device):
    dataset = DomainDataset(domains, tokenizer)
    dataloader = DataLoader(dataset, batch_size=16, shuffle=False)

    model.to(device)
    model.eval()
    predictions = []

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)

            outputs = model(input_ids, attention_mask)
            preds = torch.argmax(outputs, dim=1)
            predictions.extend(preds.cpu().numpy())

    return predictions

# 加载0527predict中的模型
def load_models():
    # 加载阶段 1 的分类器
    classifier_path = os.path.join("model", "classifier.pkl")
    with open(classifier_path, "rb") as f:
        classifier = pickle.load(f)

    # 加载阶段 2 的 APT 中心向量和 TLD 编码器
    apt_centers_path = os.path.join("model", "apt_centers.pkl")
    with open(apt_centers_path, "rb") as f:
        apt_centers, tld_encoder = pickle.load(f)

    return classifier, apt_centers, tld_encoder

# 定义每个 APT 组织的范围阈值和对应的可能性
# 这里只是示例，你可以根据实际情况调整
RANGE_THRESHOLDS = {
    "low": 0.98,
    "medium": 0.99,
    "high": 0.999
}

PROBABILITIES = {
    "low": 0.4,
    "medium": 0.6,
    "high": 0.85
}

def predict_domain(domain, classifier, apt_centers, tld_encoder):
    # 提取域名特征
    feat = extract_domain_features(domain)
    X = pd.DataFrame([feat])

    # 对 tld 列进行类别编码
    X['tld'] = X['tld'].astype('category')

    is_apt = 1
    # 使用阶段 1 的分类器进行预测
    is_apt = classifier.predict(X)[0]

    if is_apt == 0:
        return []

    # 对 TLD 进行独热编码
    tld_encoded = tld_encoder.transform(X[["tld"]]).toarray()
    tld_cols = [f"tld_{cls}" for cls in tld_encoder.categories_[0]]
    df_tld = pd.DataFrame(tld_encoded, columns=tld_cols)

    df_final = pd.concat([X.drop(columns=["tld"]).reset_index(drop=True), df_tld], axis=1)

    # 计算与每个 APT 组织中心向量的距离（使用余弦相似度）
    similarities = {}
    for apt, center in apt_centers.items():
        similarity = np.dot(df_final.values, center) / (np.linalg.norm(df_final.values) * np.linalg.norm(center))
        similarities[apt] = similarity

    # 根据相似度范围确定可能性
    probabilities = {}
    for apt, sim in similarities.items():
        if sim >= RANGE_THRESHOLDS["high"]:
            probabilities[apt] = PROBABILITIES["high"]
        elif sim >= RANGE_THRESHOLDS["medium"]:
            probabilities[apt] = PROBABILITIES["medium"]
        elif sim >= RANGE_THRESHOLDS["low"]:
            probabilities[apt] = PROBABILITIES["low"]
        else:
            probabilities[apt] = 0.0

    # 按可能性排序，选择最可能的前 3 个 APT 组织
    sorted_apt = sorted(probabilities.items(), key=lambda item: item[1], reverse=True)
    top_3_apt = sorted_apt[:3]

    return top_3_apt

# 创建 Flask 应用
app = Flask(__name__)
CORS(app) 
# 全局变量，用于存储模型和分词器
model = None
tokenizer = None
device = None
classifier = None
apt_centers = None
tld_encoder = None

# 初始化模型
# 初始化模型
def init_model():
    global model, tokenizer, device, classifier, apt_centers, tld_encoder
    try:
        logger.info("开始初始化模型...")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        tokenizer, _ = load_roberta()
        num_classes = 2
        model = DomainAPTClassifier(num_classes)

        # 加载模型参数
        checkpoint_path = os.path.join("model", "roberta_model_2025-04-22-13-30-19.pth")
        if os.path.exists(checkpoint_path):
            model.load_state_dict(torch.load(checkpoint_path, map_location=device))
            logger.info(f"成功加载模型参数: {checkpoint_path}")
        else:
            logger.error(f"模型路径不存在: {checkpoint_path}")
            raise FileNotFoundError(f"模型文件不存在: {checkpoint_path}")

        model.to(device)
        model.eval()  # 设置为评估模式
        logger.info("RoBERTa模型初始化完成")

        # 加载0527predict中的模型
        classifier, apt_centers, tld_encoder = load_models()
        logger.info("APT预测模型加载完成")
        return True
    except Exception as e:
        logger.error(f"模型初始化失败: {str(e)}")
        model = None
        return False

# 预测 API
@app.route('/predict', methods=['POST'])
def predict_api():
    if model is None:
        error_msg = "模型未初始化"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

    data = request.get_json()
    domain = data.get('domain')
    if not domain:
        error_msg = "未提供域名"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 400

    try:
        logger.info(f"收到预测请求: {domain}")
        
        # 进行预测
        predictions = predict(model, [domain], tokenizer, device)

        # 转换为标签名称
        label = "良性" if predictions[0] == 0 else "恶意"
        
        logger.info(f"域名: {domain}, 恶意性预测结果: {label}")

        result = {
            "domain": domain,
            "label": label
        }

        # 对恶意域名进行进一步预测
        if label == "恶意":
            apt_pred = predict_domain(domain, classifier, apt_centers, tld_encoder)
            result["apt_predictions"] = [{"apt": apt, "probability": prob} for apt, prob in apt_pred]
            
            logger.info(f"域名: {domain}, APT预测结果: {result['apt_predictions']}")
        else:
            result["apt_predictions"] = []
            logger.info(f"域名: {domain}, 良性域名，无需进行APT预测")

        return jsonify(result)
    except Exception as e:
        error_msg = f"预测过程中发生错误: {str(e)}"
        logger.error(error_msg, exc_info=True)  # 记录完整的错误堆栈信息
        return jsonify({"error": error_msg}), 500
    if model is None:
        return jsonify({"error": "模型未初始化"}), 500

    data = request.get_json()
    domain = data.get('domain')
    if not domain:
        return jsonify({"error": "未提供域名"}), 400

    try:
        # 进行预测
        predictions = predict(model, [domain], tokenizer, device)

        # 转换为标签名称
        label = "良性" if predictions[0] == 0 else "恶意"

        result = {
            "domain": domain,
            "label": label
        }

        # 对恶意域名进行进一步预测
        if label == "恶意":
            apt_pred = predict_domain(domain, classifier, apt_centers, tld_encoder)
            result["apt_predictions"] = [{"apt": apt, "probability": prob} for apt, prob in apt_pred]
        else:
            result["apt_predictions"] = []

        return jsonify(result)
    except Exception as e:
        return jsonify({"error": f"预测过程中发生错误: {str(e)}"}), 500

# 健康检查 API
@app.route('/health', methods=['GET'])
def health_check():
    status = {
        "status": "ok" if model is not None else "error",
        "model_initialized": model is not None,
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    return jsonify(status)

# API文档
@app.route('/', methods=['GET'])
def api_docs():
    docs = {
        " endpoints": [
            {
                "path": "/predict",
                "method": "POST",
                "description": "预测域名的恶意性和关联的APT组织",
                "request": {
                    "content-type": "application/json",
                    "body": {"domain": "example.com"}
                },
                "response": {
                    "200": {
                        "domain": "example.com",
                        "label": "恶意",
                        "apt_predictions": [
                            {"apt": "APT1", "probability": 0.85},
                            {"apt": "APT28", "probability": 0.6},
                            {"apt": "APT32", "probability": 0.4}
                        ]
                    },
                    "400": {"error": "未提供域名"},
                    "500": {"error": "模型未初始化"}
                }
            },
            {
                "path": "/health",
                "method": "GET",
                "description": "检查API的健康状态",
                "response": {
                    "200": {
                        "status": "ok",
                        "model_initialized": True,
                        "timestamp": "2025-05-27 12:00:00"
                    }
                }
            }
        ]
    }
    return jsonify(docs)

if __name__ == "__main__":
    # 初始化模型
    init_model()
    if model is not None:
        print("API服务已启动，访问 http://localhost:5000 查看文档")
        app.run(host='0.0.0.0', port=5000, debug=False)
    else:
        print("模型初始化失败，API服务无法启动")    